#include <opencv2/opencv.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/dnn.hpp>
#include <opencv2/dnn/shape_utils.hpp>
#include <image_transport/image_transport.h>

#include <ros/ros.h>
#include <signal.h>
#include <termios.h>
#include <stdio.h>
#include <cv_bridge/cv_bridge.h>
#include <sensor_msgs/Image.h>
#include <boost/thread/mutex.hpp>
#include <boost/thread/thread.hpp>

using namespace cv;
using namespace cv::dnn;
using namespace std;

const char* classNames[] = {"background",
                            "aeroplane", "bicycle", "bird", "boat",
                            "bottle", "bus", "car", "cat", "chair",
                            "cow", "diningtable", "dog", "horse",
                            "motorbike", "person", "pottedplant",
                            "sheep", "sofa", "train", "tvmonitor"};

class DetectorMobileNet
{
public:
    DetectorMobileNet(ros::NodeHandle& nh);
    ~DetectorMobileNet(){cv::destroyWindow("person_detector");}
    void start();

private:
    ros::NodeHandle& nh_;
    image_transport::ImageTransport imageTransport_;
    image_transport::TransportHints hints_;
    image_transport::Subscriber image_sub_;
    int64 work_begin_;
    double work_fps_;

    const size_t frameWidth=640;
    const size_t frameHeight = 480;
    const size_t inWidth = 300;
    const size_t inHeight = 300;
    const float WHRatio = inWidth / (float)inHeight;
    const float inScaleFactor = 0.007843f;
    const float meanVal = 127.5;
    Rect crop_;

    boost::mutex image_mutex_;
    cv::dnn::Net net_;
    cv::Mat image_;

    void workBegin();
    void workEnd();

    void imageCb(const sensor_msgs::ImageConstPtr& msg);
    void getImage(cv::Mat &image);
};



DetectorMobileNet::DetectorMobileNet(ros::NodeHandle& nh):nh_(nh),imageTransport_(nh),hints_("compressed")
{
  cv::namedWindow("person_detector");
  net_= readNetFromCaffe("MobileNetSSD_deploy.prototxt","MobileNetSSD_deploy.caffemodel");
  image_sub_ =  imageTransport_.subscribe( "/usb_cam/image_raw/compressed", 10, &DetectorMobileNet::imageCb ,this,hints_);

  Size cropSize;
  if (frameWidth / (float)frameHeight > WHRatio)
  {
      cropSize = Size(static_cast<int>(frameHeight * WHRatio),
                      frameHeight);
  }
  else
  {
      cropSize = Size(frameWidth,
                      static_cast<int>(frameWidth / WHRatio));
  }
  crop_ = Rect(Point((frameWidth - cropSize.width) / 2,
                     (frameHeight - cropSize.height) / 2),
               cropSize);
}

void DetectorMobileNet::start(){

    while(true){
       workBegin();
       cv::Mat src;
       getImage(src);
       Mat inputBlob = blobFromImage(src, inScaleFactor,
                                     Size(inWidth, inHeight), meanVal, false);

       net_.setInput(inputBlob, "data"); //set the network input

       Mat detection = net_.forward("detection_out"); //compute output
       //! [Make forward pass]

       Mat detectionMat(detection.size[2], detection.size[3], CV_32F, detection.ptr<float>());

       src = src(crop_);

       float confidenceThreshold = 0.3;
       for(int i = 0; i < detectionMat.rows; i++)
       {
           float confidence = detectionMat.at<float>(i, 2);


           if(confidence > confidenceThreshold)
           {
               size_t objectClass = (size_t)(detectionMat.at<float>(i, 1));


               if (String(classNames[objectClass])!="person"){
                   continue;
               }
               int xLeftBottom = static_cast<int>(detectionMat.at<float>(i, 3) * src.cols);
               int yLeftBottom = static_cast<int>(detectionMat.at<float>(i, 4) * src.rows);
               int xRightTop = static_cast<int>(detectionMat.at<float>(i, 5) * src.cols);
               int yRightTop = static_cast<int>(detectionMat.at<float>(i, 6) * src.rows);

               ostringstream ss;
               ss << confidence;
               String conf(ss.str());

               Rect object((int)xLeftBottom, (int)yLeftBottom,
                           (int)(xRightTop - xLeftBottom),
                           (int)(yRightTop - yLeftBottom));

               rectangle(src, object, Scalar(0, 255, 0));
               String label = String(classNames[objectClass]) + ": " + conf;
               int baseLine = 0;

               Size labelSize = getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
               rectangle(src, Rect(Point(xLeftBottom, yLeftBottom - labelSize.height),
                                     Size(labelSize.width, labelSize.height + baseLine)),
                         Scalar(255, 255, 255), CV_FILLED);
               putText(src, label, Point(xLeftBottom, yLeftBottom),
                       FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0,0,0));
           }
       }
       workEnd();

       cv::putText(src, "Person detector", cv::Point(5, 25), cv::FONT_HERSHEY_SIMPLEX, 0.8, cv::Scalar(255, 0, 0), 2);
       std::ostringstream fpsText;
       fpsText << "FPS: " << work_fps_;
       cv::putText(src, fpsText.str(), cv::Point(5, 55), cv::FONT_HERSHEY_SIMPLEX, 0.8, cv::Scalar(0, 255, 0), 2);

       cv::imshow("person_detector",src);
       cv::waitKey(1);
    }

}

void DetectorMobileNet::getImage(cv::Mat &image) {
    boost::unique_lock<boost::mutex> lk( image_mutex_ );
    image = image_.clone();
    lk.unlock();
}

void DetectorMobileNet::workBegin()
{
    work_begin_ = cv::getTickCount();
}

void DetectorMobileNet::workEnd()
{
    int64 delta = cv::getTickCount() - work_begin_;
    double freq = cv::getTickFrequency();
    work_fps_ = freq / delta;
}

void DetectorMobileNet::imageCb(const sensor_msgs::ImageConstPtr &msg)
{
    cv_bridge::CvImageConstPtr cv_ptr;

    try {
        cv_ptr = cv_bridge::toCvShare(msg, sensor_msgs::image_encodings::BGR8);
    } catch (cv_bridge::Exception& e) {
        ROS_ERROR("cv_bdrige exception: %s", e.what());
        return;
    }


    boost::unique_lock<boost::mutex> lk( image_mutex_ );
    image_=cv_ptr->image.clone();
    lk.unlock();

}

void spin() {
    ros::spin();
}

int main(int argc, char *argv[])
{

    ros::init(argc, argv, "person_detector");
    ros::NodeHandle nh;
    boost::thread spinTread(&spin);
    DetectorMobileNet detector(nh);
    detector.start();

    spinTread.join();

    return 0;
}


